from typing import List
import numpy as np
import utm
from catalog.place import MGRSRecord
from common.geometry.geom import Geom
from shapely.geometry import Point
from common.geoimage.raster_dataset import SceneMeta


def find_tileid_with_lon_lat(geom: Geom):
    tile_items: List[MGRSRecord] = MGRSRecord.query_many_items(geom)
    assert len(tile_items) > 0
    tile_items = sorted(tile_items, key=lambda x: x.uid)
    return tile_items[0].uid


def get_s2_unique_location_id(lon_lat_geom: Geom):
    tileid = find_tileid_with_lon_lat(lon_lat_geom)

    lon = lon_lat_geom._shapely_geom.x
    lat = lon_lat_geom._shapely_geom.y

    MGRSRecord.init_connection()
    mgrs_instance: MGRSRecord = MGRSRecord.get_by_id(tileid)
    MGRSRecord.close_connection()

    mgrs_meta = SceneMeta.from_dict(mgrs_instance.get_mgrs_metadict_with_resolution())
    epsg_num = mgrs_meta.prj_epsg_no
    zone_number = epsg_num % 100
    if epsg_num // 100 == 326:
        zone_letter = "N"
    elif epsg_num // 100 == 327:
        zone_letter = "S"
    else:
        raise ValueError("Not a valid epsg number.")

    affine = np.array(mgrs_meta.affine).reshape(3, 3)
    affine_inv = np.linalg.inv(affine)

    east, north, _, _ = utm.from_latlon(lat, lon, zone_number, zone_letter)
    pos_tmp = affine_inv @ np.array([[east], [north], [1]])
    col = int(np.floor(pos_tmp[0, 0]))
    row = int(np.floor(pos_tmp[1, 0]))
    return (tileid, row, col)


def get_lon_lat_from_tile_position(tileid, row, col):
    MGRSRecord.init_connection()
    mgrs_instance: MGRSRecord = MGRSRecord.get_by_id(tileid)
    MGRSRecord.close_connection()
    mgrs_meta = SceneMeta.from_dict(mgrs_instance.get_mgrs_metadict_with_resolution())
    epsg_num = mgrs_meta.prj_epsg_no
    zone_number = epsg_num % 100
    if epsg_num // 100 == 326:
        zone_letter = "N"
    elif epsg_num // 100 == 327:
        zone_letter = "S"
    else:
        raise ValueError("Not a valid epsg number.")
    affine = np.array(mgrs_meta.affine).reshape(3, 3)

    tile_pos_idx = np.array([[col + 0.5], [row + 0.5], [1]])
    pos_in_utm = affine @ tile_pos_idx
    east = pos_in_utm[0, 0]
    north = pos_in_utm[1, 0]

    lat, lon = utm.to_latlon(east, north, zone_number, zone_letter)

    return Geom.from_shapely(Point(lon, lat))


if __name__ == "__main__":
    test_location = Geom.from_shapely(Point(74.859061, 38.838656))
    unique_pos_id = get_s2_unique_location_id(test_location)
    tileid, row, col = unique_pos_id
    recovered_loc = get_lon_lat_from_tile_position(tileid, row, col)
    new_tileid, new_row, new_col = get_s2_unique_location_id(recovered_loc)
    print(tileid, new_tileid)
    print(row, new_row)
    print(col, new_col)

    print
